'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2024-09-23
Copyright © Department of Physics, Tsinghua University. All rights reserved

Access zotero data
'''

import sqlite3
import json

ZOTERO_DATA = "D:/information-base/zotero-data/zotero.sqlite"
ZOTERO_STYLE_DATA = "D:/information-base/zotero-data/storage/zot-style-storage.json"
MY_LIBRARY_ID = 1


def get_all_tags() -> list[str]:
    ''' Get all tags from current zotero file '''

    conn = sqlite3.connect(ZOTERO_DATA)
    cursor = conn.cursor()

    query = "SELECT name FROM tags"

    try:
        cursor.execute(query)

        rows = cursor.fetchall()

        all_tags = [row[0] for row in rows]

    except sqlite3.Error as e:
        print(f"An error occurred: {e.args[0]}")

    finally:
        # 关闭游标
        cursor.close()
        # 关闭数据库连接
        conn.close()

    return all_tags


def get_tags_in_library(library_id: int) -> list[str]:
    ''' Get tags in a selected library '''
    sql_query = """
        SELECT DISTINCT T.name
        FROM tags T
        JOIN itemTags IT ON T.tagID = IT.tagID
        JOIN items I ON IT.itemID = I.itemID
        WHERE I.libraryID = %d
    """ % (library_id)

    conn = sqlite3.connect(ZOTERO_DATA)
    cursor = conn.cursor()
    try:
        cursor.execute(sql_query)
        rows = cursor.fetchall()
        all_tags = [row[0] for row in rows]

    except sqlite3.Error as e:
        print(f"An error occurred: {e.args[0]}")

    finally:
        # 关闭游标
        cursor.close()
        # 关闭数据库连接
        conn.close()

    return all_tags


def parse_sharp_tags(sharp_tag: str) -> list[str]:
    ''' Parse sharp tag to a list of nested tags '''
    sharp_tag_splited = sharp_tag.split('/')
    all_tags = []
    for j in range(len(sharp_tag_splited)):
        all_tags.append(
            '/'.join(sharp_tag_splited[:len(sharp_tag_splited) - j])
        )

    return all_tags


def get_all_sharp_tags() -> set[str]:
    ''' Get all #tags from current zotero file '''
    all_tags = get_tags_in_library(MY_LIBRARY_ID)
    all_sharp_tags = []
    ordinary_tags = []
    for curr_tag in all_tags:
        if curr_tag[0] == "#":
            all_sharp_tags += parse_sharp_tags(curr_tag)
        else:
            ordinary_tags.append(curr_tag)
    return set(all_sharp_tags), set(ordinary_tags)


def tag_to_snippet(tag: str) -> str:
    ''' Get snippet string of a given tag '''
    name = "\"Tag:%s\"" % (tag)
    splitted = tag.split('/')
    if len(splitted) == 1:
        prefix = "\"%s\"" % (tag.replace("#", "").replace(" ", "_"))
    else:
        prefix = "[\"%s\", \"%s\"]" % (
            tag.replace("#", "").replace(" ", "_"),
            splitted[-1].replace(" ", "_")
        )
    body = "\"%s\"" % (tag)

    inner_str = ""
    inner_str += f"\"prefix\": {prefix}, "
    inner_str += f"\"body\": {body}, "
    inner_str += "\"scope\": \"markdown\""
    return "%s: {%s}" % (name, inner_str)


def get_all_tagged_items(library_id=MY_LIBRARY_ID) -> list[str]:
    ''' Get all item keys with tags 
        Database structure:
            items:
                itemID, ..., key
            itemTags:
                itemID, tagID
        First, traverse all "itemID" in "items", and return "key" if "itemID" in "itemTags"
    '''
    
    conn = sqlite3.connect(ZOTERO_DATA)
    cursor = conn.cursor()
    
    # SQL查询：获取所有在itemTags表中存在的items的key
    query = """
        SELECT DISTINCT I.key
        FROM items I
        JOIN itemTags IT ON I.itemID = IT.itemID
        WHERE I.libraryID = ?
    """
    
    try:
        cursor.execute(query, (library_id,))
        rows = cursor.fetchall()
        tagged_item_keys = [row[0] for row in rows]
        
    except sqlite3.Error as e:
        print(f"An error occurred: {e.args[0]}")
        tagged_item_keys = []
        
    finally:
        cursor.close()
        conn.close()
    
    return tagged_item_keys


def get_title_and_abstract(item_key: str) -> tuple[str, str]:
    ''' Get title and abstract of a given item key 
        1. Search "itemID" by "key=item_key" in table "items"
        2. Search "valueID" by "itemID" and "fieldID=1" (title) or "fieldID=2" (abstract) in table "itemData"
        3. Search "value" by "valueID" in table "itemDataValues"
    '''
    
    conn = sqlite3.connect(ZOTERO_DATA)
    cursor = conn.cursor()
    
    title = ""
    abstract = ""
    
    try:
        # Step 1: 通过key获取itemID
        cursor.execute("SELECT itemID FROM items WHERE key = ? AND libraryID = ?", (item_key, MY_LIBRARY_ID))
        item_result = cursor.fetchone()
        
        if not item_result:
            return ("", "")  # 如果找不到item，返回空字符串
        
        item_id = item_result[0]
        
        # Step 2 & 3: 获取title (fieldID=1)
        cursor.execute("""
            SELECT IDV.value 
            FROM itemData ID 
            JOIN itemDataValues IDV ON ID.valueID = IDV.valueID 
            WHERE ID.itemID = ? AND ID.fieldID = 1
        """, (item_id,))
        title_result = cursor.fetchone()
        if title_result:
            title = title_result[0]
        
        # Step 2 & 3: 获取abstract (fieldID=2)
        cursor.execute("""
            SELECT IDV.value 
            FROM itemData ID 
            JOIN itemDataValues IDV ON ID.valueID = IDV.valueID 
            WHERE ID.itemID = ? AND ID.fieldID = 2
        """, (item_id,))
        abstract_result = cursor.fetchone()
        if abstract_result:
            abstract = abstract_result[0]
            
    except sqlite3.Error as e:
        print(f"An error occurred: {e.args[0]}")
        
    finally:
        cursor.close()
        conn.close()
    
    return (title, abstract)

def get_zotero_style_data():
    with open(ZOTERO_STYLE_DATA, "r") as fp:
        style_data = json.load(fp)
    return style_data

def get_reading_time():
    data = get_zotero_style_data()
    paper_reading_time = {}
    for key, item in data.items():
        if "readingTime" in item:
            tot_time = sum([v[1] for v in item["readingTime"]["data"].items()])
            paper_reading_time[key] = tot_time
    return paper_reading_time
